#include "convolutional_layer.h"
#include "gemm.h"
#include "activations.h"
#include "im2col.h"
#include "col2im.h"

#include <math.h>

int convolutional_out_height(convolutional_layer l) {
    return (l.h + 2 * l.pad - l.size) / l.stride + 1;
}

int convolutional_out_width(convolutional_layer l) {
    return (l.w + 2 * l.pad - l.size) / l.stride + 1;
}

convolutional_layer make_convolutional_layer(int batch, int h, int w, int c, int n, int groups, 
    int size, int stride, int padding, ACTIVATION activation, int batch_normalize, int binary, int xnor, int adam) {
    int i;
    convolutional_layer l = { 0 };
    l.learning_rate_scale = 1;
    l.type = CONVOLUTIONAL;
    l.groups = groups;
    l.h = h;
    l.w = w;
    l.c = c;
    l.n = n;
    l.binary = binary;
    l.xnor = xnor;
    l.batch = batch;
    l.stride = stride;
    l.size = size;
    l.pad = padding;
    l.batch_normalize = batch_normalize;

    l.weights = calloc(c / groups * n * size * size, sizeof(float));
    l.weight_updates = calloc(c / groups * n * size * size, sizeof(float));

    l.biases = calloc(n, sizeof(float));
    l.bias_updates = calloc(n, sizeof(float));

    l.nweights = c / groups * n * size * size;
    l.nbiases = n;

    float scale = sqrt(2. / (size * size * c / l.groups));

    for (i = 0; i < l.nweights; ++i) l.weights[i] = scale * rand_normal();

    int out_w = convolutional_out_width(l);
    int out_h = convolutional_out_height(l);

    l.out_h = out_h;
    l.out_w = out_w;
    l.out_c = n;

    l.outputs = l.out_h * l.out_w * l.out_c;
    l.inputs = l.w * l.h * l.c;

    l.output = calloc(l.batch * l.outputs, sizeof(float));
    l.delta = calloc(l.batch * l.outputs, sizeof(float));

    l.forward = forward_convolutional_layer;
    l.backward = backward_convolutional_layer;
    l.update = update_convolutional_layer;

    if (l.binary) {
        l.binary_weights = calloc(l.nweights, sizeof(float)); 
        l.cweights = calloc(l.nweights, sizeof(float));
        l.scales = calloc(n, sizeof(float));
    }

    if (xnor) {
        l.binary_weights = calloc(l.nweights, sizeof(float));
        l.binary_input = calloc(l.inputs * l.batch, sizeof(float));
    }

    if (batch_normalize) {
        l.scales = calloc(n, sizeof(float));
        l.scale_updates = calloc(n, sizeof(float));
        for (i = 0; i < n; i++) {
            l.scales[i] = 1;
        }

        l.mean = calloc(n, sizeof(float));
        l.variance = calloc(n, sizeof(float));

        l.mean_delta = calloc(n, sizeof(float));
        l.variance_delta = calloc(n, sizeof(float));

        l.rolling_mean = calloc(n, sizeof(float));
        l.rolling_variance = calloc(n, sizeof(float));

        l.x = calloc(l.batch * l.outputs, sizeof(float));
        l.x_norm = calloc(l.batch * l.outputs, sizeof(float));
    }

    if (adam) {
        l.m = calloc(l.nweights, sizeof(float));
        l.v = calloc(l.nweights, sizeof(float));

        l.bias_m = calloc(n, sizeof(float));
        l.scale_m = calloc(n, sizeof(float));

        l.bias_v = calloc(n, sizeof(float));
        l.scale_v = calloc(n, sizeof(float));
    }

#ifdef GPU
    l.forward_gpu = forward_convolutional_layer_gpu;
    l.backward_gpu = backward_convolutional_layer_gpu;
    l.update_gpu = update_convolutional_layer_gpu;

    if (gpu_index >= 0) {
        if (adam) {
            l.m_gpu = cuda_make_array(l.m, l.nweights);
            l.v_gpu = cuda_make_array(l.v, l.nweights);
            l.bias_m_gpu = cuda_make_array(l.bias_m, n);
            l.bias_v_gpu = cuda_make_array(l.bias_v, n);
            l.scale_m_gpu = cuda_make_array(l.scale_m, n);
            l.scale_v_gpu = cuda_make_array(l.scale_v, n);
        }

        l.weights_gpu = cuda_make_array(l.weights, l.nweights);
        l.weight_updates_gpu = cuda_make_array(l.weight_updates, l.nweights);

        l.biases_gpu = cuda_make_array(l.biases, n);
        l.bias_updates_gpu = cuda_make_array(l.bias_updates, l.nweights);

        l.delta_gpu = cuda_make_array(l.delta, l.batch * out_h * out_w * n);
        l.output_gpu = cuda_make_array(l.output, l.batch * out_h * out_w * n);

        if (binary) {
            l.binary_weights_gpu = cuda_make_array(l.weights, l.nweights);
        }

        if (xnor) {
            l.binary_weights_gpu = cuda_make_array(l.weights, l.nweights);
            l.binary_input_gpu = cuda_make_array(0, l.inputs * l.batch);
        }

        if (batch_normalize) {
            l.mean_gpu = cuda_make_array(l.mean, n);
            l.variance_gpu = cuda_make_array(l.variance, n);

            l.rolling_mean_gpu = cuda_make_array(l.mean, n);
            l.rolling_variance_gpu = cuda_make_array(l.variance, n);

            l.mean_delta_gpu = cuda_make_array(l.mean, n);
            l.variance_delta_gpu = cuda_make_array(l.variance, n);

            l.scales_gpu = cuda_make_array(l.scales, n);
            l.scale_updates_gpu = cuda_make_array(l.scale_updates, n);

            l.x_gpu = cuda_make_array(l.output, l.batch * out_h * out_w * n);
            l.x_norm_gpu = cuda_make_array(l.output, l.batch * out_h * out_w * n);
        }

#ifdef CUDNN
        cudnnCreateTensorDescriptor(&l.normTensorDesc);
        cudnnCreateTensorDescriptor(&l.srcTensorDesc);
        cudnnCreateTensorDescriptor(&l.dstTensorDesc);
        cudnnCreateFilterDescriptor(&l.weightDesc);
        cudnnCreateTensorDescriptor(&l.dsrcTensorDesc);
        cudnnCreateTensorDescriptor(&l.ddstTensorDesc);
        cudnnCreateTensorDescriptor(&l.dweightDesc);
        cudnnCreateConvolutionDescriptor(&l.convDesc);
        cudnn_convolutional_setup(&l);
#endif // CUDNN
    }

#endif
    l.workspace_size = get_workspace_size(l);
    l.activation = activation;

    fprintf(stderr, "conv  %5d %2d x%2d /%2d  %4d x%4d x%4d   ->  %4d x%4d x%4d  %5.3f BFLOPs\n", n, size, size, stride, w, h, c, l.out_w, l.out_h, l.out_c, (2.0 * l.n * l.size*l.size*l.c/l.groups * l.out_h*l.out_w) / 1000000000.);
    return l;
}

void forward_convolutional_layer(convolutional_layer l, network net) {
    int i, j;
    fill_cpu(l.outputs * l.batch, 0, l.output, 1);

    if (l.xnor) {
        
    }

    int m = l.n / l.groups;
    int k = l.size * l.size * l.c / l.groups;
    int n = l.out_w * l.out_h;
    for (i = 0; i < l.batch; ++i) {
        for (j = 0; j < l.groups; ++j) {
            float *a = l.weights + j * l.nweights / l.groups;
            float *b = net.workspace;
            float *c = l.output + (i * l.groups + j) * n * m;
            float *im = net.input + (i * l.groups + j) * l.c / l.groups * l.h * l.w;

            if (l.size = 1) {
                b = im;
            } else {
                im2col_cpu(im, l.c / l.groups, l.h, l.w, l.size, l.stride, l.pad, b);
            }
            gemm(0, 0, m, n, k, 1, a, k, b, n, 1, c, n);
        }
    }

    if (l.batch_normalize) {

    } else {
        add_bias(l.output, l.biases, l.batch, l.n, l.out_h * l.out_w);
    }

    // activate_array(); 
}

void backward_convolutional_layer(convolutional_layer l, network net) {
    int i, j;
    int m = l.n / l.groups;
    int n = l.size * l.size * l.c / l.groups;
    int k = l.out_w * l.out_h;

    gradient_array(l.output, l.outputs * l.batch, l.activation, l.delta);

    if (l.batch_normalize) {

    } else {
        backward_bias(l.bias_updates, l.delta, l.batch, l.n, k);
    }

    for (i = 0; i < l.batch; ++i) {
        for (j = 0; j < l.groups; ++j) {
            float *a = l.delta + (i * l.groups + j) * m * k;
            float *b = net.workspace;
            float *c = l.weight_updates + j * l.nweights / l.groups;

            float *im = net.input + (i * l.groups + j) * l.c / l.groups * l.h * l.w;
            float *imd = net.delta + (i * l.groups + j) * l.c / l.groups * l.h * l.w;
            
            if (l.size == 1) {
                b = im;
            } else {
                //im2col_cpu(im, l.c / l.groups, l.h, l.w, l.size, l.stride, l.pad, b);
            }

            gemm(0, 1, m, n, k, 1, a, k, b, k, 1, c, n);

            if (net.delta) {
                a = l.weights + j * l.nweights / l.groups;
                b = l.delta + (i * l.groups + j) * m * k;
                c = net.workspace;
                
                if (l.size == 1) {
                    c = imd;
                }

                gemm(1, 0, n, k, m, 1, a, n, b, k, 0, c, k);

                if (l.size != 1) {
                    col2im_cpu(net.workspace, l.c / l.groups, l.h, l.w, l.size, l.stride, l.pad, imd);
                }
            }
        }
    }
}

void update_convolutional_layer(convolutional_layer l, update_args arg) {
    float learning_rate = arg.learning_rate * l.learning_rate_scale;
    float momentum = arg.momentum;
    float decay = arg.decay;
    int batch = arg.batch;
    
    axpy_cpu(l.n, learning_rate / batch, l.bias_updates, 1, l.biases, 1);
    scal_cpu(l.n, momentum, l.bias_updates, 1);
    
    if (l.scales) {

    }

    axpy_cpu(l.nweights, -decay * batch, l.weights, 1, l.weight_updates, 1);
    axpy_cpu(l.nweights, learning_rate / batch, l.weight_updates, 1, l.weights, 1);
    scal_cpu(l.nweights, momentum, l.weight_updates, 1);
}

void update_convolutional_adam_layer(convolutional_layer l, update_args arg) {
    // 更新一阶矩
    axpy_cpu(l.nweights, (1 - arg.beta1), l.weight_updates, 1, l.m, 1);
    scal_cpu(l.nweights, arg.beta1, l.weight_updates, 1);

    // 更新二阶矩
    mul_cpu(l.nweights, l.weight_updates, 1, l.weight_updates, 1);
    axpy_cpu(l.nweights, (1 - arg.beta2), l.weight_updates, 1, l.v, 1);
    scal_cpu(l.nweights, arg.beta2, l.weight_updates, 1);

    // 修正偏差
    /*
    float m_hat = l.m[j] / (1 - powf(beta1, net.batches_done + 1));
    float v_hat = l.v[j] / (1 - powf(beta2, net.batches_done + 1));
    */
    // 更新参数
    //l.weights[j] -= learning_rate * m_hat / (sqrtf(v_hat) + epsilon);
}

void add_bias(float *output, float *biases, int batch, int n, int size) {
    int i, j, b;
    for (b = 0; b < batch; ++b) {
        for (i = 0; i < n; ++i) {
            for (j = 0; j < size; ++j) {
                output[(b * n + i) * size + j] += biases[i];
            }
        }
    }
}